package com.me.test.controller;

import com.me.test.config.WebSocketConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * websocket：
 * 1，客户端通过 @ServerEndpoint(value = "/test/one") 指定的 URI 连接服务端
 * 2，指定 configurator 才能使用 Spring 的依赖注入，否则注入 null
 *
 * @author ME
 * @date 2022/3/4
 */
@Slf4j
@ServerEndpoint(value = "/ws/{sid}", configurator = WebSocketConfig.CustomSpringConfigurator.class)
@Component
public class WebSocketApi {

    private static final String HEARTBEAT = "heartbeat";

    public WebSocketApi() {
        System.out.println("websocket init");
    }

    /** 记录当前在线连接数 */
    private static final AtomicInteger ONLINE_COUNT = new AtomicInteger(0);

    /** 存放所有在线的客户端 */
    private static final Map<String, Session> CLIENTS = new ConcurrentHashMap<>();

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session, @PathParam(value = "sid") String sid) {
        ONLINE_COUNT.incrementAndGet(); // 在线数加1
        CLIENTS.put(sid, session);
        log.info("有新连接加入：{}，当前在线人数为：{}", sid, ONLINE_COUNT.get());

        this.sendMessage("Hello, <span style=\"color:#FF0000\">" + sid + "</span>", sid, session);
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose(@PathParam(value = "sid") String sid) {
        ONLINE_COUNT.decrementAndGet(); // 在线数减1
        CLIENTS.remove(sid);
        log.info("有一连接关闭：{}，当前在线人数为：{}", sid, ONLINE_COUNT.get());
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message, Session session, @PathParam(value = "sid") String sid) {
        // 处理心跳消息
        if (HEARTBEAT.equals(message)) {
            return;
        }

        if (log.isDebugEnabled()) {
            log.debug("收到 [{}] 的消息：{}", sid, message);
        }
        // 回复给消息发送者
        // this.sendMessage("Hello, " + message, sid, session);
        // 群发消息
        this.massTexting(message, sid, session, false);
    }

    @OnError
    public void onError(Session session, Throwable error, @PathParam(value = "sid") String sid) {
        log.error("[" +sid + "]发生错误", error);
    }

    /**
     * 服务端发送消息给客户端
     */
    private void sendMessage(String message, String sid, Session toSession) {
        try {
            if (log.isDebugEnabled()) {
                log.debug("向 [{}] 发送消息：{}", sid, message);
            }
            toSession.getBasicRemote().sendText(message);
        } catch (Exception e) {
            log.error("发送消息失败：", e);
        }
    }

    /**
     * 群发消息
     *
     * @param message 消息内容
     * @param includeSelf 群发是否包括自己
     */
    private void massTexting(String message, String sid, Session fromSession, boolean includeSelf) {
        if (includeSelf) {
            for (Map.Entry<String, Session> sessionEntry : CLIENTS.entrySet()) {
                if (log.isDebugEnabled()) {
                    log.debug("向 [{}] 发送消息：{}", sessionEntry.getKey(), message);
                }
                sessionEntry.getValue().getAsyncRemote().sendText(message);
            }
        } else {
            Session session = CLIENTS.remove(sid);
            for (Map.Entry<String, Session> sessionEntry : CLIENTS.entrySet()) {
                if (log.isDebugEnabled()) {
                    log.debug("向 [{}] 发送消息：{}", sessionEntry.getKey(), message);
                }
                sessionEntry.getValue().getAsyncRemote().sendText(message);
            }
            if (session != null) {
                CLIENTS.put(sid, session);
            }
        }
    }

}
